/*
 *  \file cvmfs_fsck.cc
 *  This tool checks a cvmfs2 cache directory for consistency.
 *  If necessary, the managed cache db is removed so that 
 *  it will be rebuilt on next mount.
 */
 
#define _FILE_OFFSET_BITS 64
#include "config.h"
#include "util.h"
#include "hash.h"
#include "atomic.h"

extern "C" {
   #include "compression.h"
}

#include <iostream>
#include <string>
#include <cstring>
#include <pthread.h>
#include <stdint.h>
#include <stdlib.h>
#include <unistd.h>
#include <dirent.h>
#include <stdio.h>
#include <errno.h>
#include <sys/stat.h>

using namespace std;

static void usage() {
   cout << "This tool checks a cvmfs2 cache directory for consistency." << endl;
   cout << "If necessary, the managed cache db is removed so that" << endl;
   cout << "it will be rebuilt on next mount." << endl << endl;
   cout << "Usage: cvmfs_fsck [-p] [-f] [-j #threads] <cache directory>" << endl;
   cout << "Options:" << endl;
   cout << "  -p try to fix automatically" << endl;
   cout << "  -f force rebuild of managed cache db on next mount" << endl;
   cout << "  -j number of concurrent integrity check worker threads" << endl;
}


struct dir_data {
   DIR *dp; ///< opened directory pointer
   int dir_num; ///< between 0 and 255
};
typedef struct dir_data dir_data_t;

string cache_dir;
atomic_int num_files;
pthread_mutex_t mutex_output = PTHREAD_MUTEX_INITIALIZER;
int concurrent_threads_max = 4;
int concurrent_threads = 0;
pthread_mutex_t mutex_ct = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t cond_ct = PTHREAD_COND_INITIALIZER;

bool fix_errors = false;
bool force_rebuild = false;
pthread_mutex_t mutex_force_rebuild = PTHREAD_MUTEX_INITIALIZER;

static void *check_dir(void *start_data) {
   dir_data_t *dir_data = (dir_data_t *)start_data;
   
   /* simple work queue: don't have more than a certain number of concurrent threads */
   pthread_mutex_lock(&mutex_ct);
   do {
      pthread_cond_wait(&cond_ct, &mutex_ct);
   } while (concurrent_threads > concurrent_threads_max);
   concurrent_threads++;
   pthread_mutex_unlock(&mutex_ct);
   
   struct dirent64 *d;
   while ((d = readdir64(dir_data->dp)) != NULL) {
      const string name = d->d_name;
      if ((name == ".") || (name == "..")) continue;
      char hex[3];
      snprintf(hex, 3, "%02x", dir_data->dir_num);
      const string lpath = string(hex) + "/" + name;
      const string path = cache_dir + "/" + lpath;
      const string hash_name = string(hex) + name;
      
      int n = atomic_xadd(&num_files, 1);
      if ((n % 1000) == 0) {
         pthread_mutex_lock(&mutex_output);
         cout << "." << flush;
         pthread_mutex_unlock(&mutex_output);
      }
         
      if (d->d_type != DT_REG) {
         pthread_mutex_lock(&mutex_output);
         cout << endl << "Warning: " << path << " is not a regular file" << endl;
         pthread_mutex_unlock(&mutex_output);
         continue;
      }
         
      /* Compress every file into /tmp and check sha1 key */
      FILE *fsrc = fopen(lpath.c_str() , "r");
      if (!fsrc) {
         pthread_mutex_lock(&mutex_output);
         cout << endl << "Error: cannot open " << path << endl;
         pthread_mutex_unlock(&mutex_output);
         continue;
      }
         
      hash::t_sha1 sha1;
      if (compress_file_sha1_only(fsrc, sha1.digest) != 0) {
         pthread_mutex_lock(&mutex_output);
         cout << endl << "Error: could not compress " << path << endl;
         pthread_mutex_unlock(&mutex_output);
      } else {
         if (sha1.to_string() != hash_name) {
            if (fix_errors) {
               if (unlink(lpath.c_str()) == 0) {
                  pthread_mutex_lock(&mutex_output);
                  cout << endl << "Fix: " << path << " is corrupted, file unlinked" << endl;
                  pthread_mutex_unlock(&mutex_output);
                  
                  /* Changes made, we have to rebuild the managed cache db */
                  pthread_mutex_lock(&mutex_force_rebuild);
                  force_rebuild = true;
                  pthread_mutex_unlock(&mutex_force_rebuild);
               } else {
                  pthread_mutex_lock(&mutex_output);
                  cout << endl << "Error: " << path << " is corrupted, could not unlink" << endl;
                  pthread_mutex_unlock(&mutex_output);
               }
            } else {
               pthread_mutex_lock(&mutex_output);
               cout << endl << "Error: " << path << " has compressed checksum " << sha1.to_string() << 
                              ", delete this file from cache directory!" << endl;
               pthread_mutex_unlock(&mutex_output);
            }
         }
      }
      fclose(fsrc);         
   }
   
   /* this one goes down, wake up another thread */
   pthread_mutex_lock(&mutex_ct);
   concurrent_threads--;
   pthread_mutex_unlock(&mutex_ct);
   pthread_cond_signal(&cond_ct);
   return NULL;
}

int main(int argc, char **argv) {
   char c;
   while ((c = getopt(argc, argv, "hpfj:")) != -1) {
      switch (c) {
         case 'h':
            usage();
            return 0;
         case 'p':
            fix_errors = true;
            break;
         case 'f':
            force_rebuild = true;
            break;
         case 'j':
            concurrent_threads_max = atoi(optarg);
            if (concurrent_threads_max < 1) {
               cout << "There is at least one worker thread required" << endl;
               return 2;
            }
            break;
         case '?':
         default:
            abort();
      }
   }
     
   /* Switch to cache directory */
   if (optind >= argc) {
      usage();
      return 2;
   }
   cache_dir = canonical_path(argv[optind]);
   if (chdir(cache_dir.c_str()) != 0) {
      cerr << "could not change to " << cache_dir << endl;
      return 1;
   }
   
   /* Check if txn directory is empty */
   DIR *dirp_txn;
   if ((dirp_txn = opendir("txn")) == NULL) {
      cerr << "invalid cache directory, " << cache_dir << "/txn does not exist" << endl;
      return 1;
   }
   struct dirent64 *d;
   while ((d = readdir64(dirp_txn)) != NULL) {
      const string name = d->d_name;
      if ((name == ".") || (name == "..")) continue;
      
      if (!fix_errors) {
         cout << "Warning: temporary directory " << cache_dir << "/txn is not empty" << endl;
         break;
      }
      const string path = "txn/" + name;
      if (unlink(path.c_str()) == 0) {
         cout << "Fix: unlinking temporary file " << path << endl;
      } else {
         cout << "Error: could not unlink " << path << " (" << errno << ")" << endl;
      }
   }
   closedir(dirp_txn);
   
   /* Open all subdirectories */
   cout << "Checking file integrity: " << flush;
   DIR *dirp[256];
   int total_files = 0;
   for (int i = 0; i <= 0xff; ++i) {
      char hex[3];
      snprintf(hex, 3, "%02x", i);
      if ((dirp[i] = opendir(hex)) == NULL) {
         cerr << "invalid cache directory, " << cache_dir << "/" << hex << " does not exist" << endl;
         return 1;
      }
      
      struct dirent64 *d;
      while ((d = readdir64(dirp[i])) != NULL) {
         const string name = d->d_name;
         if ((name == ".") || (name == "..")) continue;
         total_files++;
      }
      rewinddir(dirp[i]);
   }
   
   /* Check file integrity, multithreaded */
   cout << total_files << " files found " << flush;
   atomic_init(&num_files);
   pthread_t worker[256];
   dir_data_t dir_data[256];
   for (int i = 0; i <= 0xff; ++i) {
      dir_data[i].dp = dirp[i];
      dir_data[i].dir_num = i;
      if (pthread_create(&worker[i], NULL, check_dir, (void *)(&dir_data[i])) != 0) {
         cout << "Fatal: could not create worker thread" << endl;
         return 1;
      }
   }
   pthread_cond_broadcast(&cond_ct);
   for (int i = 0; i <= 0xff; ++i) {
      pthread_join(worker[i], NULL);
      closedir(dirp[i]);
   }
   cout << endl;
   
   if (atomic_read(&num_files) != total_files) {
      cout << "Warning: cache directory has been touched during check (" << 
              atomic_read(&num_files) << " checked)" << endl;
   }
   
   if (force_rebuild) {
      if (unlink("cvmfscatalog.cache") == 0) {
         cout << "Fix: managed cache db unlinked, will be rebuilt on next mount" << endl;
      } else {
         if (errno != ENOENT) {
            cout << "Error: could not unlink managed cache database (" << errno << ")" << endl;
         }
      }
   }
   
   /* Check cached catalogs */
   /*cout << "Checking cached catalogs: " << flush;
   DIR *dirp_clg;
   if ((dirp_clg = opendir(".")) == NULL) {
      cerr << "unable to open cache directory" << endl;
      return 1;
   }
   while ((d = readdir64(dirp_clg)) != NULL) {
      const string name = d->d_name;
      
      if (name.substr(0, 15) == "cvmfs.checksum.") {
         cout << "." << flush;
         const string name_clg = "cvmfs.catalog." + name.substr(15);
         struct stat64 info;
         if (stat64(name_clg.c_str(), &info) != 0) {
            if (!fix_errors) {
               cout << "Warning: abandoned checksum " << cache_dir << "/" << name << 
                       ", delete that file" << endl;
            } else {
               if (unlink(name.c_str()) == 0) {
                  cout << "Fix: unlinking abandoned checksum file " << name << endl;
               } else {
                  cout << "Error: could not unlink " << name << " (" << errno << ")" << endl;
               }
            }
            continue;
         } 
         
         hash::t_sha1 sha1_clg;
         hash::t_sha1 sha1_chksum;
         FILE *fchksum = fopen(name.c_str(), "r");
         char tmp[40];
         if (fchksum && (fread(tmp, 1, 40, fchksum) == 40)) {
            sha1_chksum.from_hash_str(string(tmp, 40));
         } else {
            cout << "Error: cannot open " << name << endl;
         }
         if (fchksum) fclose(fchksum);
         FILE *fsrc = fopen(name_clg.c_str() , "r");
         if (!fsrc) {
            cout << endl << "Error: cannot open " << name_clg << endl;
            continue;
         }
         if (compress_file_sha1_only(fsrc, sha1_clg.digest) != 0) {
            cout << endl << "Error: could not compress " << name_clg << endl;
         }
         if (sha1_clg != sha1_chksum) {
            if (!fix_errors) {
               cout << "Warning: cached checksum does not match catalog, " 
                    << cache_dir << "/" << name << " and " << cache_dir << "/" << name_clg << 
                       ", delete these files" << endl;
            } else {
               if ((unlink(name.c_str())) == 0 && (unlink(name_clg.c_str()) == 0)) {
                  cout << "Fix: unlinking unmatching checksum and catalog files " << name << ", " << name_clg << endl;
               } else {
                  cout << "Error: could not unlink " << name << ", " << name_clg << " (" << errno << ")" << endl;
               }
            }
            continue;
         }
      }
      
      if (name.substr(0, 14) == "cvmfs.catalog.") {
         struct stat64 info;
         if (stat64(("cvmfs.checksum." + name.substr(14)).c_str(), &info) != 0) {
            if (!fix_errors) {
               cout << "Warning: abandoned catalog " << cache_dir << "/" << name << 
                       ", delete that file" << endl;
            } else {
               if (unlink(name.c_str()) == 0) {
                  cout << "Fix: unlinking abandoned catalog file " << name << endl;
               } else {
                  cout << "Error: could not unlink " << name << " (" << errno << ")" << endl;
               }
            }
         }
      }
   }
   closedir(dirp_clg);
   cout << endl;*/
   
   return 0;
}
